package com.zrc.japktool.apk.jiagu;

import com.zrc.japktool.util.Zip4jUtil;
import net.lingala.zip4j.ZipFile;
import net.lingala.zip4j.model.FileHeader;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

public class JGFileUtil {
    // 加密密钥，两边必须保持一致
    private static final byte KEY = 0x5A;

    /**
     * 加密数据
     */
    private static byte[] encrypt(byte[] data) {
        if (data == null) return null;

        byte[] result = new byte[data.length];
        for (int i = 0; i < data.length; i++) {
            result[i] = (byte) (data[i] ^ KEY);
        }
        return result;
    }

    /**
     * 解密数据
     */
    private static byte[] decrypt(byte[] data) {
        // XOR加密解密使用相同算法
        return encrypt(data);
    }

    /**
     * 写入JG文件
     * @param filePath 输出文件路径
     * @param packageName 包名
     * @param dexList dex文件数据列表
     * @throws IOException IO异常
     */
    public static void createJGFile(String filePath, String packageName, List<byte[]> dexList) throws IOException {
        // 验证参数合法性
        if (packageName == null || packageName.isEmpty()) {
            throw new IllegalArgumentException("packageName cannot be null or empty");
        }
        if (dexList == null || dexList.isEmpty()) {
            throw new IllegalArgumentException("dexList cannot be empty");
        }

        try (FileOutputStream fos = new FileOutputStream(filePath); DataOutputStream dos = new DataOutputStream(fos)) {

            // 写入包名长度和包名（包名不加密）
            byte[] packageBytes = packageName.getBytes(StandardCharsets.UTF_8);
            dos.writeInt(packageBytes.length);
            dos.write(packageBytes);

            // 写入dex文件数量
            dos.writeInt(dexList.size());

            // 写入每个dex文件（加密）
            for (int i = 0; i < dexList.size(); i++) {
                byte[] dexData = dexList.get(i);
                if (dexData == null || dexData.length == 0) {
                    throw new IllegalArgumentException("dexList contains empty or null data at index: " + i);
                }

                byte[] encryptedDex = encrypt(dexData);
                dos.writeInt(encryptedDex.length);
                dos.write(encryptedDex);
            }
        }
    }

    /**
     * 生成JG文件的字节数组
     * @param packageName 包名
     * @param dexList dex文件数据列表
     * @return JG文件的字节数组
     * @throws IOException IO异常
     */
    public static byte[] createJGFile(String packageName, List<byte[]> dexList) throws IOException {
        // 验证参数合法性
        if (packageName == null || packageName.isEmpty()) {
            throw new IllegalArgumentException("packageName cannot be null or empty");
        }
        if (dexList == null || dexList.isEmpty()) {
            throw new IllegalArgumentException("dexList cannot be null or empty");
        }

        // 使用 ByteArrayOutputStream 在内存中构建文件内容
        try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); DataOutputStream dos = new DataOutputStream(baos)) {

            // 写入包名（使用UTF-8编码确保兼容性）
            byte[] packageBytes = packageName.getBytes(StandardCharsets.UTF_8);
            dos.writeInt(packageBytes.length);
            dos.write(packageBytes);

            // 写入dex文件数量
            dos.writeInt(dexList.size());

            // 写入每个加密后的dex文件
            for (int i = 0; i < dexList.size(); i++) {
                byte[] dexData = dexList.get(i);
                if (dexData == null || dexData.length == 0) {
                    throw new IllegalArgumentException("dexList contains empty or null data at index: " + i);
                }

                byte[] encryptedDex = encrypt(dexData);
                dos.writeInt(encryptedDex.length);
                dos.write(encryptedDex);
            }

            // 刷新流并返回完整字节数组
            dos.flush();
            return baos.toByteArray();
        }
    }

    /**
     * 读取JG文件
     * @param filePath 文件路径
     * @return 包含包名和dex列表的结果对象
     * @throws IOException IO异常
     */
    public static JGFileData parseJGFile(String filePath) throws IOException {
        // 验证文件是否存在
        File file = new File(filePath);
        if (!file.exists() || !file.isFile()) {
            throw new FileNotFoundException("JG file not found: " + filePath);
        }

        try (FileInputStream fis = new FileInputStream(filePath)) {
            return parseJGFile(fis);
        }
    }

    /**
     * 从字节数组解析JG文件
     * @param fileBytes JG文件的字节数组
     * @return 包含包名和dex列表的结果对象
     * @throws IOException IO异常
     */
    public static JGFileData parseJGFile(byte[] fileBytes) throws IOException {
        // 验证输入参数
        if (fileBytes == null || fileBytes.length == 0) {
            throw new IllegalArgumentException("fileBytes cannot be null or empty");
        }

        // 使用ByteArrayInputStream从字节数组读取数据
        try (ByteArrayInputStream bais = new ByteArrayInputStream(fileBytes)) {
            return parseJGFile(bais);
        }
    }

    /**
     * 从输入流读取JG文件
     * @param inputStream 输入流
     * @return 包含包名和dex列表的结果对象
     * @throws IOException IO异常
     */
    public static JGFileData parseJGFile(InputStream inputStream) throws IOException {
        if (inputStream == null) {
            throw new IllegalArgumentException("inputStream cannot be null");
        }

        // 使用try-with-resources确保流自动关闭
        try (DataInputStream dis = new DataInputStream(inputStream)) {

            // 读取包名
            int packageLen = dis.readInt();
            if (packageLen <= 0) {
                throw new IOException("Invalid package name length: " + packageLen);
            }

            byte[] packageBytes = new byte[packageLen];
            int bytesRead = dis.read(packageBytes);
            if (bytesRead != packageLen) {
                throw new IOException("Could not read complete package name. Expected: "
                        + packageLen + ", read: " + bytesRead);
            }
            String packageName = new String(packageBytes, StandardCharsets.UTF_8);

            // 读取dex文件数量
            int dexCount = dis.readInt();
            if (dexCount < 0) {
                throw new IOException("Invalid dex count: " + dexCount);
            }

            List<byte[]> dexList = new ArrayList<>(dexCount);

            // 读取每个dex文件
            for (int i = 0; i < dexCount; i++) {
                int dexLen = dis.readInt();
                if (dexLen <= 0) {
                    throw new IOException("Invalid dex length at index " + i + ": " + dexLen);
                }

                byte[] encryptedDex = new byte[dexLen];
                bytesRead = dis.read(encryptedDex);
                if (bytesRead != dexLen) {
                    throw new IOException("Could not read complete dex data at index " + i
                            + ". Expected: " + dexLen + ", read: " + bytesRead);
                }

                byte[] decryptedDex = decrypt(encryptedDex);
                dexList.add(decryptedDex);
            }

            // 验证是否有未预期的数据
            if (dis.available() > 0) {
                throw new IOException("Unexpected remaining data in JG file: " + dis.available() + " bytes");
            }

            return new JGFileData(packageName, dexList);
        }
    }

    // 用于存储读取结果的数据类
    public static class JGFileData {
        private final String packageName;
        private final List<byte[]> dexList;

        public JGFileData(String packageName, List<byte[]> dexList) {
            this.packageName = packageName;
            this.dexList = dexList;
        }

        public String getPackageName() {
            return packageName;
        }

        public List<byte[]> getDexList() {
            return dexList;
        }
    }

    public static List<byte[]> getApkDexBytes(String apkPath) {
        List<byte[]> dexBytesList = new ArrayList<>();
        List<FileHeader> headersToRemove = new ArrayList<>();

        try (ZipFile zipFile = new ZipFile(apkPath)) {
            List<FileHeader> fileHeaders = zipFile.getFileHeaders();

            for (FileHeader fileHeader : fileHeaders) {
                String fileName = fileHeader.getFileName();
                // 筛选根目录下的Dex文件
                if (fileName.endsWith(".dex") && !fileName.endsWith("/") && !fileName.contains("/")) {
                    // 使用缓冲流可靠读取文件内容
                    try (InputStream inputStream = zipFile.getInputStream(fileHeader);
                         ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {

                        byte[] buffer = new byte[8192];
                        int bytesRead;
                        while ((bytesRead = inputStream.read(buffer)) != -1) {
                            outputStream.write(buffer, 0, bytesRead);
                        }
                        byte[] dexData = outputStream.toByteArray();

                        if (dexData.length > 0) {
                            dexBytesList.add(dexData);
                            headersToRemove.add(fileHeader);
                        }
                    }
                }
            }

            // 删除标记的Dex文件
            for (FileHeader header : headersToRemove) {
                zipFile.removeFile(header);
            }
        } catch (Exception e) {
            System.err.println("读取APK中Dex文件时发生错误: " + e.getMessage());
            e.printStackTrace();
        }
        return dexBytesList;
    }

    // 演示用法
    public static void main(String[] args) throws IOException {
        String apkPath = "C:/Users/15666/Desktop/test/C_jg-zipalign-sign.apk";
        byte [] libjgBytes = Zip4jUtil.readFileBytes(apkPath, "static/lib.jg");
        JGFileData jg = JGFileUtil.parseJGFile(libjgBytes);
        String outApkPath = "C:/Users/15666/Desktop/test/D.apk";
        for (int i = 0; i < jg.getDexList().size(); i++) {
            Zip4jUtil.writeFileBytes(outApkPath,"classes"+i+".dex",jg.getDexList().get(i),false);
        }
    }
}
